# incoding=utf-8
import socket
import struct
import hashlib
import threading,random
import sys
import base64
reload(sys)
sys.setdefaultencoding('utf-8')

connectionlist = {}

def sendMessage(message):
    global connectionlist
    for connection in connectionlist.values():
        tmp = base64.b64encode(message)
        bstr = send_data(tmp)
        print message
        b = connection.send(bstr)

def deleteconnection(item):
    global connectionlist
    del connectionlist['connection'+item]

#接收客户端发送过来的消息,并且解包
def RecvData(nNum,client):
    try:
        pData = client.recv(nNum)
        if not len(pData):
            return False
    except:
        return False
    else:
        code_length = ord(pData[1]) & 127
        if code_length == 126:
            masks = pData[4:8]
            data = pData[8:]
        elif code_length == 127:
            masks = pData[10:14]
            data = pData[14:]
        else:
            masks = pData[2:6]
            data = pData[6:]
        
        raw_str = ""
        i = 0
        for d in data:
            raw_str += chr(ord(d) ^ ord(masks[i%4]))
            i += 1
        return raw_str
        
#打包发送数据给客户端
def SendData(pData,client):
    if(pData == False):
        return False
    else:
        pData = str(pData)
        
    token = "\x81"
    length = len(pData)
    if length < 126:
        token += struct.pack("B", length)
    elif length <= 0xFFFF:
        token += struct.pack("!BH", 126, length)
    else:
        token += struct.pack("!BQ", 127, length)
    pData = '%s%s' % (token,pData)

    client.send(pData)
    
    return True

def send_data(raw_str):
    back_str = []

    back_str.append('\x81')
    data_length = len(raw_str)

    if data_length < 125:
        back_str.append(chr(data_length))
    else:
        back_str.append(chr(126))
        back_str.append(chr(data_length >> 8))
        back_str.append(chr(data_length & 0xFF))

    back_str = "".join(back_str) + raw_str    
    return back_str
     
class WebSocket(threading.Thread):
    def __init__(self,conn,index,name,remote, path="/"):
        threading.Thread.__init__(self)
        self.conn = conn
        self.index = index
        self.name = name
        self.remote = remote
        self.path = path
        self.buffer = ""
        
    def run(self):
        print 'Socket%s Start!' % self.index
        headers = {}
        self.handshaken = False

        while True:
            if self.handshaken == False:
                print 'Socket%s Start Handshaken with %s!' % (self.index,self.remote)
                self.buffer += self.conn.recv(1024)
                if self.buffer.find('\r\n\r\n') != -1:
                    header, data = self.buffer.split('\r\n\r\n', 1)
                    for line in header.split("\r\n")[1:]:
                        key, value = line.split(": ", 1)
                        headers[key] = value
                    print 'header:-->'+header
                    headers["Location"] = "ws://%s%s" %(headers["Host"], self.path)
                    self.buffer = data[8:]
                    key = headers["Sec-WebSocket-Key"]
                    token = self.generate_token(key)
                    
                    handshake = '\
HTTP/1.1 101 Web Socket Protocol Handshake\r\n\
Upgrade: webSocket\r\n\
Connection: Upgrade\r\n\
Sec-WebSocket-Accept:%s\r\n\
Sec-WebSocket-Origin: %s\r\n\
Sec-WebSocket-Location: %s\r\n\r\n\
' %(token,headers['Origin'], headers['Location'])

                    print handshake
                    num = self.conn.send(handshake)
                    print str(num)
                    self.handshaken = True
                    print 'Socket%s Handshaken with %s success!' % (self.index,self.remote)
                    #self.conn.send('\x00'+'Welcome'+'\xFF'+'\n')
                    welcomeData = "Welcome 你的ID是:"+ str(self.index)
                    welcomeData = base64.b64encode(welcomeData)
                    bstr = send_data(welcomeData)
                    sendMessage('系统消息:ID为*'+str(self.index)+'*的用户加入聊天室')
                    print welcomeData
                    self.conn.send(bstr)
                    #SendData('\x00'+'Welcome'+'\xFF',self.conn)
                    #sendMessage('Welcome, '+self.name+' !')
            else:
                self.buffer = RecvData(8196,self.conn)                
                if self.buffer:
                    print 'rec:'+self.buffer                   
                if self.buffer:
                    #s = self.buffer.split("\xFF")[0][1:]
                    t = self.buffer
                    s = base64.b64decode(t)
                    if s=='quit':
                        print 'Socket%s Logout!' % (self.index)
                        sendMessage('系统消息:ID为*'+str(self.index)+'*的用户退出聊天室')
                        deleteconnection(str(self.index))
                        self.conn.close()
                        break
                    else:
                        print 'Socket%s Got msg:%s from %s!' % (self.index,t,self.remote)
                        print 'msg decode:%s' % (s)
                        sendMessage(str(self.index)+':'+self.name+':'+s)
                    self.buffer = ""
    
    def generate_token(self, key):
        import base64
        nkey=key+'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
        nkey=base64.b64encode(hashlib.sha1(nkey).digest())
        return nkey

class WebSocketServer(object):
    def __init__(self):
        self.socket = None
    def begin(self):
        print 'WebSocketServer Start!'
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.bind(("127.0.0.1",8000))
        self.socket.listen(50)
        
        global connectionlist
        
        i=0
        while True:
            connection, address = self.socket.accept()
            username=address[0]
            newSocket = WebSocket(connection,i,username,address)
            newSocket.start()
            connectionlist['connection'+str(i)]=connection
            print str(connection)
            i = i + 1

if __name__ == "__main__":
    server = WebSocketServer()
    server.begin()